import hfscf
import numpy as np

def SCF (r = 1.4632, Z=[1,1], b1 = hfscf.GTO["H"], b2 = hfscf.GTO["H"], b = 3, vbs=True):
    # Задаются центры атомов
    R = [0, r]
    
    # генерируется матрица перекрывания
    if vbs: print("*) Generando matriz de traslape S.")
    s_scf = hfscf.S(R, b1, b2, b)
    
    # Генерируется гамильтониан H для оператора Фока
    if vbs: print("\n*) Generando hamiltoniano H.")
    h_scf = hfscf.H(R, Z, b1, b2, b)
        
    # Диагонализация матрицы S и поиск диагональной матрицы X (X = S^(-1/2))
    # Xa - сопряженно-транспонированная матрица
    if vbs: print("\n*) Diagonalizando matriz S y hallando matriz diagonal X.")
    X = hfscf.diagon(m=s_scf)
    Xa = X.getH()
    
    # Создание матрицы плотности P
    if vbs: print("\n*) Creando matriz de densidad P.")
    p_scf = np.matrix([[0,0],[0,0]], dtype=np.float64)  # Referencia (7) p. 148
    
    # Проверка сходимости SCF (итеративно заполняем матрицу F, диагонализируем, проверяем разницу)
    if vbs: print("\n*) Comenzando con el SCF.")
    for iteracion in range(50):
        # Расчет двух электронной части матрицы Фока
        # F = H + G
        if vbs: print("\n**) Generando la matriz de Fock: calculando \
integrales de dos electrones.")
        g_scf = hfscf.G(r, p_scf, b1, b2, b)
        f_scf = h_scf + g_scf  # Referencia (7) p. 141 eq. (3.154)
        
        # Создание матрицы F' (меняем базис F)
        # F' = X_adj * F * X
        if vbs: print("**) Cambiando la base de F.")
        f_tra = Xa * f_scf * X
        
        # Диагонализация F' и генерация C'
        if vbs: print("**) Diagonalizando F' y generando C'.")
        c_tra = hfscf.diagon2(m=f_tra)
        
        # Построение матрицы коэффициентов C
        # C = X * C'
        if vbs: print("**) Construyendo matriz de coeficientes C.")
        c_scf = X * c_tra
        
        # Из матрицы C строим матрицу Р
        # Пересчет матрицы плотности P
        if vbs: print("**) Recalculando matriz de densidad P.")
        p_temp = hfscf.P(C=c_scf)
        
        print("\nConcluida la " + str(iteracion + 1) + ". iteracion.\n")
        
        # Проверка сходимости
        if np.linalg.norm(p_temp - p_scf) < 1E-4: # Referencia (7) p. 148
            print("\n\n-->El campo autoconsistente SI ha convergido!")
            return {"S":s_scf,"H":h_scf,"X": X,"F":f_scf,"C":c_scf,"P":p_temp}
        else:
            p_scf = p_temp
    print("\n\n-->El campo autoconsistente NO ha convergido!\nRevisar supuestos.")
    return {"S":s_scf,"H":h_scf,"X": X,"F":f_scf,"C":c_scf,"P":p_temp}

orbs = SCF()
print(orbs)

*) Generando matriz de traslape S.

*) Generando hamiltoniano H.

*) Diagonalizando matriz S y hallando matriz diagonal X.

*) Creando matriz de densidad P.

*) Comenzando con el SCF.

**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 1. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 2. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 3. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 4. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 5. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 6. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 7. iteracion.



-->El campo autoconsistente SI ha convergido!
{'S': matrix([[0.99999999, 0.63749012],
        [0.63749012, 0.99999999]]), 'H': matrix([[-1.09920375, -0.91954841],
        [-0.91954841, -1.09920375]]), 'X': matrix([[ 0.55258063,  1.17442445],
        [ 0.55258063, -1.17442445]]), 'F': matrix([[-0.34684107, -0.57771139],
        [-0.57771139, -0.34684028]]), 'C': matrix([[ 0.55258114, -1.17442421],
        [ 0.55258013,  1.17442468]]), 'P': matrix([[0.61069182, 0.61069071],
        [0.61069071, 0.6106896 ]])}

Орбитали в 1D

hfscf.orbital(orbs['C'])

Орбитали в 2D (слева связывающая, справа разрыхляющая)

hfscf.orbital2D(orbs['C'], orbs['X'], orbs['F'])

orben = hfscf.ener_orbs(orbs['X'], orbs['F'])
elecen = hfscf.ener_elec(orbs['P'], orbs['H'], orbs['F'])
toten = hfscf.ener_tot()
print(f'Энергии орбиталей: {orben[0]} и {orben[1]}')
print(f'Электронная энергия молекулы: {elecen}')
print(f'Полная энергия молекулы: {toten}')

Энергии орбиталей: -0.5646153594583119 и 0.6368673980935621
Электронная энергия молекулы: -1.7974485548087507
Полная энергия молекулы: 0.683433570256971